#include "stinger-atomics.h"
#include "stinger-utils.h"
#include "stinger.h"
#include "xmalloc.h"
#include "bfs.h"
#include "mathUtils.c"

#define INC 1
int getCountK33(const struct stinger *G, uint64_t v, uint64_t *visited, uint64_t i, uint64_t j, uint64_t k, uint64_t nv);
int getCountK33(const struct stinger *G, uint64_t v, uint64_t *visited, uint64_t i, uint64_t j, uint64_t k, uint64_t nv)
{
  int countK33 = 0;
  int k33found = 0;
  uint64_t candidates = 0;
  uint64_t * nodes = (uint64_t *)xcalloc(nv, sizeof(uint64_t));
  size_t deg_i, deg_j, deg_k, d;
  deg_i = stinger_outdegree (G, i);
  deg_j = stinger_outdegree (G, j);
  deg_k = stinger_outdegree (G, k);

  uint64_t * neighbors_i = (uint64_t *)xmalloc(sizeof(uint64_t) * deg_i);
  uint64_t * neighbors_j = (uint64_t *)xmalloc(sizeof(uint64_t) * deg_j);
  uint64_t * neighbors_k = (uint64_t *)xmalloc(sizeof(uint64_t) * deg_k);
  uint64_t * k33candidates = (uint64_t *)xmalloc(sizeof(uint64_t) * deg_k);

  stinger_gather_typed_successors (G, 0, i, &d, neighbors_i, deg_i);
  assert(d == deg_i);
  stinger_gather_typed_successors (G, 0, j, &d, neighbors_j, deg_j);
  assert(d == deg_j);
  stinger_gather_typed_successors (G, 0, k, &d, neighbors_k, deg_k);
  assert(d == deg_k);
 
  for(size_t iter = 0 ; iter < deg_i ; iter++)
    nodes[neighbors_i[iter]] = 1;

  for(size_t iter = 0 ; iter < deg_j ; iter++)
  {
    if(nodes[neighbors_j[iter]] == 1)
      nodes[neighbors_j[iter]] = 2;
  }

  for(size_t iter = 0 ; iter < deg_k ; iter++)
  {
    if(nodes[neighbors_k[iter]] == 2)
      nodes[neighbors_k[iter]] = 3;
  }

  for(uint64_t iter = 0 ; iter < nv ; iter++)
  {
    if(nodes[iter] == 3 && stinger_has_typed_successor(G, 0, v, iter) == 0)
    {
      k33candidates[candidates] = iter;
      candidates++;
    }
      countK33++;
  }
  
  for(i = 0; i < candidates; i++)
  {
    for(j =  i + 1; j < candidates; j++)
    {
      if(stinger_has_typed_successor(G, 0, k33candidates[i], k33candidates[j]) == 0)
      {
	visited[k33candidates[i]] = 1;
	visited[k33candidates[j]] = 1;
	visited[v] = 1;
	k33found = 1;
	break;
      }
    }
  }

  free(k33candidates);
  free(neighbors_i);
  free(neighbors_j);
  free(neighbors_k);
  free(nodes);

  return k33found;
}

/* This function is returns the K3,1 count for a node */
int findK31CountForNode(const struct stinger *G , uint64_t v, uint64_t nv, uint64_t *visited){

      size_t deg_v = stinger_outdegree (G, v);
      int64_t w_start = 0;	
      int64_t my_w = stinger_int64_fetch_add (&w_start, deg_v);
      size_t d;
      int countK31 = 0;
      int i, j, k;	
      int countK33 = 0;
      int k33found = 0;

      /* Find index of all neighbors of node v */
      uint64_t * neighbors = (uint64_t *)xmalloc(sizeof(uint64_t) * deg_v);
      stinger_gather_typed_successors (G, 0, v, &d, neighbors, deg_v);
      assert(d == deg_v);

      for (i = 0 ; i < deg_v; i++){
	if(visited[neighbors[i]] == 0) {

	  for (j = i + 1 ; j < deg_v; j++){

	    if(stinger_has_typed_successor(G, 0, neighbors[i], neighbors[j]) == 0
	       && visited[neighbors[j]] == 0)
	    {
	      
	      for(k = j + 1 ; k < deg_v; k++){
		/* Check if any two pairs have an edge, if not we obtain a K31 pair */
		if(visited[neighbors[k]] == 0
		   && stinger_has_typed_successor(G, 0, neighbors[j], neighbors[k]) == 0 
		   && stinger_has_typed_successor(G, 0, neighbors[i], neighbors[k]) == 0)
		{
		  k33found = 0;
		  k33found = getCountK33(G, v, visited, neighbors[i], neighbors[j], neighbors[k], nv);
		  if(k33found)
		  {
		    visited[neighbors[i]] = 1;
		    visited[neighbors[j]] = 1;
		    visited[neighbors[k]] = 1;
		    countK33++;
		  }
		}
	      }
	    }
	  }
	}
      } 
      free(neighbors);
      return countK33; 
}
/* This function is used to set up the environment for starting BFS */
int setBFS(const struct stinger *G, const int64_t nv, const int64_t ne, const int64_t s, const int64_t numSteps , int *countK31InitGraph){

  int64_t k, x;
  int64_t *Q_big = (int64_t *) xmalloc (INC * nv * sizeof (int64_t));
  int64_t *marks_s_big = (int64_t *) xmalloc (INC * nv * sizeof (int64_t));
  int64_t *QHead_big = (int64_t *) xmalloc (INC * 2 * numSteps * sizeof (int64_t));
  int64_t *neighbors_big = (int64_t *) xmalloc (INC * ne * sizeof (int64_t));

  int64_t count = 0;

  k = 0;
  x = 0;
  
  OMP ("omp parallel for")
    MTA ("mta assert parallel")
    MTA ("mta loop future")
    MTA ("mta assert nodep")
    MTA ("mta assert no alias")
    for (x = 0; x < INC; x++) 
    {
      int64_t *Q = Q_big + x * nv;
      int64_t *marks_s = marks_s_big + x * nv;
      int64_t *QHead = QHead_big + x * 2 * numSteps;
      int64_t *neighbors = neighbors_big + x * ne;

      *countK31InitGraph = bfs_stinger (G, nv, ne, s, marks_s, numSteps, Q, QHead, neighbors);
      printf("Number of 3,3s: %d\n", *countK31InitGraph);
    }

  free (neighbors_big);
  free (QHead_big);
  free (marks_s_big);
  free (Q_big);
  return *countK31InitGraph;
}

int bfs_stinger (const struct stinger *G, const int64_t nv, const int64_t ne,
             const int64_t startV,
             int64_t * marks, const int64_t numSteps,
             int64_t * Q, int64_t * QHead, int64_t * neighbors)
{
  /* Variable Declaration */
  int countK31InitGraph = 0;
  int64_t j, k, s;
  int64_t nQ, Qnext, Qstart, Qend;
  int64_t w_start;
  int64_t d_phase;
  int64_t *visited = (int64_t *) xcalloc (nv, sizeof (int64_t));

  MTA ("mta assert nodep")
    for (j = 0; j < nv; j++) {
      marks[j] = 0;
    }

  s = startV;
  /* Push node s onto Q and set bounds for first Q sublist */
  Q[0] = s;
  Qnext = 1;
  nQ = 1;
  QHead[0] = 0;
  QHead[1] = 1;
  marks[s] = 1;
    
  PushOnStack:                   /* Push nodes onto Q */
 
  /* Execute the nested loop for each node v on the Q AND for each neighbor w of v  */
  d_phase = nQ;
  Qstart = QHead[nQ - 1];
  Qend = QHead[nQ];
  w_start = 0;
  countK31InitGraph = findK31CountForNode(G , s, nv, visited);

  
  MTA ("mta assert no dependence")
    MTA ("mta block dynamic schedule")
    for (j = Qstart; j < Qend; j++) {
      int64_t v = Q[j];
      size_t d;
      size_t deg_v = stinger_outdegree (G, v);
      int64_t my_w = stinger_int64_fetch_add (&w_start, deg_v);
      stinger_gather_typed_successors (G, 0, v, &d, &neighbors[my_w], deg_v);
      assert (d == deg_v);

      MTA ("mta assert nodep")
        for (k = 0; k < deg_v; k++) {
          int64_t d, w, l;
          w = neighbors[my_w + k];

          /* If node has not been visited, set distance and push on Q */
          if (stinger_int64_fetch_add (marks + w, 1) == 0) {

	    /* Check if this node forms a K3,1 */
	    if (stinger_outdegree (G, w) >= 3){	
	      if(visited[v] == 0)
		countK31InitGraph += findK31CountForNode(G , w, nv, visited);
	    }	
            Q[stinger_int64_fetch_add (&Qnext, 1)] = w;
          }
        }
    }


  if (Qnext != QHead[nQ] && nQ < numSteps) {
    nQ++;
    QHead[nQ] = Qnext;
    goto PushOnStack;
  }

  free(visited);
  return countK31InitGraph;
}

int64_t 
st_conn_stinger (const struct stinger *G, const int64_t nv, const int64_t ne,
                 const int64_t * sources, const int64_t num,
                 const int64_t numSteps)
{
  int64_t k, x;

  int64_t *Q_big = (int64_t *) xmalloc (INC * nv * sizeof (int64_t));
  int64_t *marks_s_big = (int64_t *) xmalloc (INC * nv * sizeof (int64_t));
  int64_t *marks_t_big = (int64_t *) xmalloc (INC * nv * sizeof (int64_t));
  int64_t *QHead_big =
    (int64_t *) xmalloc (INC * 2 * numSteps * sizeof (int64_t));
  int64_t *neighbors_big = (int64_t *) xmalloc (INC * ne * sizeof (int64_t));

  int64_t count = 0;

  k = 0;
  x = 0;

  OMP ("omp parallel for")
    MTA ("mta assert parallel")
    MTA ("mta loop future")
    MTA ("mta assert nodep")
    MTA ("mta assert no alias")
    for (x = 0; x < INC; x++) {
      int64_t *Q = Q_big + x * nv;
      int64_t *marks_s = marks_s_big + x * nv;
      int64_t *marks_t = marks_t_big + x * nv;
      int64_t *QHead = QHead_big + x * 2 * numSteps;
      int64_t *neighbors = neighbors_big + x * ne;

      for (int64_t claimedk = stinger_int64_fetch_add (&k, 2);
           claimedk < 2 * num; claimedk = stinger_int64_fetch_add (&k, 2)) {
        int64_t s = sources[claimedk];
        int64_t deg_s = stinger_outdegree (G, s);
        int64_t t = sources[claimedk + 1];
        int64_t deg_t = stinger_outdegree (G, t);

        if (deg_s == 0 || deg_t == 0) {
          stinger_int64_fetch_add (&count, 1);
        } else {
          bfs_stinger (G, nv, ne, s, marks_s, numSteps, Q, QHead, neighbors);
          bfs_stinger (G, nv, ne, t, marks_t, numSteps, Q, QHead, neighbors);
          int64_t local_count = 0;

          MTA ("mta assert nodep")
            for (int64_t j = 0; j < nv; j++) {
              if (marks_s[j] && marks_t[j])
                stinger_int64_fetch_add (&local_count, 1);
            }

          if (local_count == 0)
            stinger_int64_fetch_add (&count, 1);
        }
      }
    }

  free (neighbors_big);
  free (QHead_big);
  free (marks_t_big);
  free (marks_s_big);
  free (Q_big);

  return count;
}


int64_t
st_conn_stinger_source (const struct stinger * G, const int64_t nv,
                        const int64_t ne, const int64_t from,
                        const int64_t * sources, const int64_t num,
                        const int64_t numSteps)
{
  int64_t k;

  int64_t *Q = (int64_t *) xmalloc (nv * sizeof (int64_t));
  int64_t *marks_s = (int64_t *) xmalloc (nv * sizeof (int64_t));
  int64_t *marks_t = (int64_t *) xmalloc (nv * sizeof (int64_t));
  int64_t *QHead = (int64_t *) xmalloc (2 * numSteps * sizeof (int64_t));
  int64_t *neighbors = (int64_t *) xmalloc (ne * sizeof (int64_t));

  int64_t count = 0;

  int64_t deg_s = stinger_outdegree (G, from);
  if (deg_s == 0) {
    return num;
  }

  bfs_stinger (G, nv, ne, from, marks_s, numSteps, Q, QHead, neighbors);

  for (k = 0; k < num; k++) {
    int64_t t = sources[k];
    int64_t deg_t = stinger_outdegree (G, t);

    if (deg_t == 0) {
      stinger_int64_fetch_add (&count, 1);
    } else {
      bfs_stinger (G, nv, ne, t, marks_t, numSteps, Q, QHead, neighbors);

      int64_t local_count = 0;

      MTA ("mta assert nodep")
        for (int64_t j = 0; j < nv; j++) {
          if (marks_s[j] && marks_t[j])
            stinger_int64_fetch_add (&local_count, 1);
        }

      if (local_count == 0)
        stinger_int64_fetch_add (&count, 1);
    }
  }

  free (neighbors);
  free (QHead);
  free (marks_t);
  free (marks_s);
  free (Q);

  return count;

}


void
bfs_csr (const int64_t nv, const int64_t ne, const int64_t * off,
         const int64_t * ind, const int64_t startV, int64_t * marks,
         const int64_t numSteps, int64_t * Q, int64_t * QHead)
{
  int64_t j, k, s;
  int64_t nQ, Qnext, Qstart, Qend;

  int64_t d_phase;

  MTA ("mta assert nodep")
    for (j = 0; j < nv; j++) {
      marks[j] = 0;
    }

  s = startV;
  /* Push node s onto Q and set bounds for first Q sublist */
  Q[0] = s;
  Qnext = 1;
  nQ = 1;
  QHead[0] = 0;
  QHead[1] = 1;
  marks[s] = 1;

 PushOnStack:                   /* Push nodes onto Q */

  /* Execute the nested loop for each node v on the Q AND
     for each neighbor w of v  */
  d_phase = nQ;
  Qstart = QHead[nQ - 1];
  Qend = QHead[nQ];

  MTA ("mta assert no dependence")
    MTA ("mta block dynamic schedule")
    for (j = Qstart; j < Qend; j++) {
      int64_t v = Q[j];
      int64_t myStart = off[v];
      int64_t myEnd = off[v + 1];

      MTA ("mta assert nodep")
        for (k = myStart; k < myEnd; k++) {
          int64_t d, w, l;
          w = ind[k];
          /* If node has not been visited, set distance and push on Q */
          if (stinger_int64_fetch_add (marks + w, 1) == 0) {
            Q[stinger_int64_fetch_add (&Qnext, 1)] = w;
          }
        }
    }

  if (Qnext != QHead[nQ] && nQ < numSteps) {
    nQ++;
    QHead[nQ] = Qnext;
    goto PushOnStack;
  }
}


int64_t
st_conn_csr (const int64_t nv, const int64_t ne, const int64_t * off,
             const int64_t * ind, const int64_t * sources, const int64_t num,
             const int64_t numSteps)
{
  int64_t k, x;

  int64_t *Q_big = (int64_t *) xmalloc (INC * nv * sizeof (int64_t));
  int64_t *marks_s_big = (int64_t *) xmalloc (INC * nv * sizeof (int64_t));
  int64_t *marks_t_big = (int64_t *) xmalloc (INC * nv * sizeof (int64_t));
  int64_t *QHead_big =
    (int64_t *) xmalloc (INC * 2 * numSteps * sizeof (int64_t));
  int64_t *neighbors_big = (int64_t *) xmalloc (INC * ne * sizeof (int64_t));

  int64_t count = 0;

  k = 0;

  MTA ("mta assert parallel")
    MTA ("mta loop future")
    for (x = 0; x < INC; x++) {
      int64_t *Q = Q_big + x * nv;
      int64_t *marks_s = marks_s_big + x * nv;
      int64_t *marks_t = marks_t_big + x * nv;
      int64_t *QHead = QHead_big + x * 2 * numSteps;
      int64_t *neighbors = neighbors_big + x * ne;

      for (int64_t claimedk = stinger_int64_fetch_add (&k, 2);
           claimedk < 2 * num; claimedk = stinger_int64_fetch_add (&k, 2)) {

        int64_t s = sources[claimedk];
        int64_t t = sources[claimedk + 1];
        bfs_csr (nv, ne, off, ind, s, marks_s, numSteps, Q, QHead);
        bfs_csr (nv, ne, off, ind, t, marks_t, numSteps, Q, QHead);

        int64_t local_count = 0;

        MTA ("mta assert nodep")
          for (int64_t j = 0; j < nv; j++) {
            if (marks_s[j] && marks_t[j])
              stinger_int64_fetch_add (&local_count, 1);
          }

        if (local_count == 0)
          stinger_int64_fetch_add (&count, 1);
      }
    }

  free (neighbors_big);
  free (QHead_big);
  free (marks_t_big);
  free (marks_s_big);
  free (Q_big);

  return count;
}
